
import os

import numpy as np
import tensorflow as tf
from tensorflow import keras

"""
神经网络训练
"""



prefix = ".npy"
# 训练集
dataPath = "D:/img/tmpIdcard/data"
# 标签集
valuePath = "D:/img/tmpIdcard/value"
# 临时存储
checkpointSavePath = "d:/Ai/checkpoint/idcard6.ckpt"
# 模型保存路径
modelPath = "d:/model/idcardLearn"


x_train = np.load(dataPath + prefix)
y_train = np.load(valuePath + prefix)
# 归一化
x_train = x_train / 255.0

# 随机打乱数据
np.random.seed(226)
np.random.shuffle(x_train)
np.random.seed(226)
np.random.shuffle(y_train)
tf.random.set_seed(226)

model = keras.Sequential()

# 卷积层 要求输入图片维度，
model.add(keras.layers.
          Conv2D(filters=128, kernel_size=(3, 3), strides=1, padding='same', activation='relu'))
model.add(keras.layers.BatchNormalization())
model.add(keras.layers.MaxPool2D(pool_size=(2, 2), strides=1, padding='same'))

# 全连接层
model.add(keras.layers.Flatten())
model.add(keras.layers.Dense(128, activation="relu"))
# model.add(keras.layers.Dense(64, activation="relu"))
model.add(keras.layers.Dropout(0.2))
model.add(keras.layers.Dense(11, activation="softmax", kernel_regularizer=tf.keras.regularizers.l2()))

model.compile(optimizer='adam',
              loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=False),
              metrics=['sparse_categorical_accuracy'])


if os.path.exists(checkpointSavePath + '.index'):
    print('-------------load the model-----------------')
    model.load_weights(checkpointSavePath)

cp_callback = tf.keras.callbacks.ModelCheckpoint(filepath=checkpointSavePath,
                                                 save_weights_only=True,
                                                 save_best_only=True)

history = model.fit(x_train, y_train, epochs=20,
                    validation_split=0.3,
                    validation_freq=5,
                    callbacks=[cp_callback]
                    )

model.summary()
model.save(filepath=modelPath)
# cv2.waitKey(0)
# # 删除建立的全部窗口
# cv2.destroyAllWindows()

val = model.predict(x_train)
val = tf.nn.softmax(val)
pred = tf.argmax(val, axis=1)
for index in range(len(pred)):
    print("pred == {} value = {}".format(pred[index], y_train[index]))
